import pandas as pd
import numpy as np
from Bio.SeqUtils.ProtParam import ProteinAnalysis
import panel as pn
import hvplot.pandas
import holoviews as hv
import matplotlib.pyplot as plt
from holoviews import opts
from scipy import optimize
import warnings
warnings.filterwarnings('ignore')

evidences = pd.read_csv('../output/evidence_aligned.csv')

# Amino acids favoring secondary structures (Levitt, M. Biochemistry 17, 4277–4285 (1978))
evidences['helix_fraction'] = [(seq.count('A') + seq.count('L') + seq.count('M') + seq.count('H') + 
                                seq.count('Q') + seq.count('E'))/len(seq)  for seq in evidences['Sequence']]
evidences['sheet_fraction'] = [(seq.count('V') + seq.count('I') + seq.count('F') + seq.count('T') + 
                                seq.count('Y'))/len(seq) for seq in evidences['Sequence']]
evidences['turn_fraction'] = [(seq.count('G') + seq.count('S') + seq.count('D') + seq.count('N') + 
                               seq.count('P'))/len(seq) for seq in evidences['Sequence']]

evidences['lastAA'] = evidences['Sequence'].str[-1:]
evidences_trp = evidences[evidences.lastAA.isin(['K', 'R'])]

# Supplementary Figure 5
sec_str = ['helix_fraction', 'sheet_fraction', 'turn_fraction']

select_sec_str = pn.widgets.Select(name='Select a secondary structure', 
                                   options=sec_str,
                                   width=340
                                  )

@pn.depends(select_sec_str.param.value)
def visualize_sec_str(sec_str):
    
    scatter_plot = evidences_trp.hvplot.points(
        x='m/z', y='CCS', c=sec_str, xlabel='m/z', ylabel='CCS (\u212b\u00B2)', clabel='Fraction', 
        width=450, height=400, cmap=plt.get_cmap('YlOrRd_r'), tools=['hover'], 
        colorbar=True, rasterize=True)
    
    scatter_plot.opts(
        toolbar='above', clim=(0, 0.61)
    )
    
    return scatter_plot 

layout = pn.Column(
    pn.pane.Markdown(
        """### Supplementary Figure 5. Fraction of amino acids favoring helical (A, L, M, H, Q, E), turn (V, I, F, T, Y) and sheet (G, S, D, N, P) secondary structures according to Levitt 1978.""", 
        margin=(10, 0, 20, 0)
    ),
    pn.Row(
        pn.layout.VSpacer(width=110),
        select_sec_str, 
        visualize_sec_str
    ),
    align='center',
    width=970
)
layout.embed()